#ifndef M_SPPM
#define M_SPPM
#include<list>
#include<omp.h>
#include"raytracer.h"
#include"../accelerators/Kd4Pts.h"
using namespace std;
#define MAXDEPTH 8
#define INITRADIUS 0.5f
const Float GAMMA = 2. / 3.;



//Progressive Photon Mapping
class SPPM : public RayTracer
{
private:
	inline unsigned int hash(const int ix, const int iy, const int iz)
	{
		return (unsigned int)(
			(ix * 73856093) ^
			(iy * 19349663) ^
			(iz * 83492791)) % hash_grid.size();
	}
	void build_hash_grid(const int w, const int h)
	{
		// determine hash table size
		// we now find the bounding box of all the measurement points inflated by the initial radius
		hpbox.reset();
		Float maxR = 0;
		int photon_count = 0;
		for (auto itr = visible_pts.begin(); itr != visible_pts.end(); ++itr)
		{
			VisiblePt *hp = (*itr);
			if (hp->beta.IsBlack()) continue;
			Float r = sqrtf(film[hp->id].r2);
			photon_count++;
			hpbox.merge(hp->p - r);
			hpbox.merge(hp->p + r);
			maxR = max(maxR, r);
		}

		P3 diagonal = hpbox.maxi - hpbox.mini;
		Float maxDiag = diagonal[hpbox.maxExtent()];
		int baseGridRes = (int)(maxDiag / maxR);
		for (int i = 0; i < 3; ++i)
			gridRes[i] = max((int)(baseGridRes * diagonal[i] / maxDiag), 1);

		for (int i=0;i<visible_pts.size();++i)
		{

			VisiblePt *hp =visible_pts[i];
			Float r = sqrtf(film[hp->id].r2);
			int imin[3], imax[3];
			toGrid(hp->p - r, imin);
			toGrid(hp->p + r, imax);
			for (int iz = imin[2]; iz <= imax[2]; iz++)
			{
				for (int iy = imin[1]; iy <= imax[1]; iy++)
				{
					for (int ix = imin[0]; ix <= imax[0]; ix++)
					{
						int hv = hash(ix, iy, iz);

						hash_grid[hv].push_back(hp);
					}
				}
			}
		}

	}
	void gen_camera_path(long long seed);
	void gen_photon_path(int n_iter, int n_sample);
	void save_img(char* name, int iter, int n_sample);
	//locate the hash grid with neibouring points
	inline bool toGrid(P3 &x, int ansPt[3])
	{
		bool inBound = true;
		P3 offset = x - hpbox.mini;
		for (int i = 0; i < 3; ++i)
		{
			offset[i] /= (hpbox.maxi[i] - hpbox.mini[i]);
			ansPt[i] = (int)gridRes[i] * offset[i];
			inBound &= (ansPt[i] >= 0 && ansPt[i] < gridRes[i]);
			ansPt[i] = clamp(ansPt[i], 0, gridRes[i] - 1);
		}
		return inBound;
	}
	int maxDepth;
	Float initRadius;
	int gridRes[3];
	BoundingBox hpbox;
	vector<VisiblePt*> visible_pts;
	sppmkd::KdTreeAccel *kd_tree;
	vector<vector<VisiblePt*>> hash_grid;
	string outputname;
	SPPMPixel* film = NULL;
	int n_pixels;
public:
	SPPM(bool enableMedia) :RayTracer(RT_SPPM,enableMedia), maxDepth(MAXDEPTH), initRadius(INITRADIUS) { hpbox.reset(); }
	~SPPM() {}
	void render(string outputname, string config);

};
void SPPM::gen_camera_path(long long seed)
{
	auto start = chrono::system_clock::now();
	int w = cam->w, h = cam->h;
	int num_procs = omp_get_max_threads();
	list<VisiblePt*>* vispts = new list<VisiblePt*>[num_procs]; //for parrellel computation

	//generate visible points for all pixels
#pragma omp parallel for schedule(dynamic, 1)
	for (int y = 0; y < h; y++)
	{
		int thread_id = omp_get_thread_num();
		for (int x = 0; x < w; x++)
		{
			int index = w * (h-1-y) + (w-1-x);
			long long cur_seed = seed + x*x+y*y;
			HaltonGenerator randGen(cur_seed);
			Float r1 = 2 * randGen.genRand(), dx = r1<1 ? sqrt(r1) - 1 : 1 - sqrt(2 - r1);
			Float r2 = 2 * randGen.genRand(), dy = r2<1 ? sqrt(r2) - 1 : 1 - sqrt(2 - r2);
			Float tx = (Float)(x + dx / 2) / w, ty = (Float)(y + dy / 2) / h;
			Ray ray = cam->gen_ray(tx, ty, NULL);
			Spectrum beta(1);
			bool isSpec = false; //whether the reflection is specular
			for (int depth = 0; depth < maxDepth; ++depth)
			{
				Float t;
				SurfaceInteraction si;
				bool sampleMedia = false;
				P3 its_pt = ray.o;
				bool intersect = scene->nearest_intersect(ray, &t, &si);
				if (enableMedia&&ray.medium)
				{
					Float sampled_t;
					beta *= ray.medium->sampleMediumInteraction(ray, t, &sampled_t, &randGen);
					if (beta.IsBlack()) break;
					sampleMedia = sampled_t < t;
					its_pt = ray.get(sampled_t);
				}
				if (sampleMedia)
				{
					if (intersect)
					{
						Object* obj = static_cast<Object*>(si.pObj);
						P3 surface_pt = ray.get(t);
						Material* material = obj->material;
						shared_ptr<BSDF> bsdf;
						if (material)
							bsdf = material->createBSDF(&si, TransportMode::Radiance);
						//choose light
						Float u0, lightPdf;
						u0 = randGen.genRand();
						int lightId = min((int)(u0*scene->lights.size()), scene->lights.size() - 1);
						lightPdf = 1.f / (scene->lights.size());
						Light* light = scene->lights[lightId];
						Ray isect_ray = { its_pt,ray.d,ray.medium };
						film[index].Ld += beta * sampleDirectLighting(&si, bsdf, isect_ray, light, &randGen, true) / lightPdf;
					}
					Float pdf;
					Float rands[] = { randGen.genRand(),randGen.genRand() };
					P3 nextDir = ray.medium->sample_inScatterDir(its_pt, -ray.d, &pdf, rands);
					isSpec = false;
					ray.o = its_pt;
					ray.d = nextDir;
					if (si.pObj) {
						MediumInterface mi = static_cast<Object*>(si.pObj)->mi;
						ray.medium = mi.getMedium(si.n, nextDir);
					}
					
					continue;
				}
				if (!intersect)
				{
					for (auto const &light : scene->lights)
					{
						film[index].Ld += beta*light->lightEnv(ray);
					}
					break;
				}
				Object* obj = static_cast<Object*>(si.pObj);
				Material* material = obj->material;
				its_pt = ray.get(t);
				int into = 0;
				P3 nl = si.n.dot(ray.d) < 0 ? into = 1, si.n : si.n * -1;
				
				if (!material)  //intersecting boundary of media
				{
					ray.o = ray.get(t + eps);
					MediumInterface mi = static_cast<Object*>(si.pObj)->mi;
					ray.medium = mi.getMedium(si.n, ray.d);
					//--depth;
					continue;
				}
				shared_ptr<BSDF> bsdf = material->createBSDF(&si, TransportMode::Radiance);
				//compute direct illum
				if (depth == 0 || isSpec)
				{
					//to be implemented: direct illum from obj emission
					film[index].Ld += beta * obj->lightOut(&si,-ray.d);
					
				}
				//accumulate direct illum from lights
				Float rands[] = { randGen.genRand(),randGen.genRand() };

				//choose one light to sample based on light power distribution
				Float u0, lightPdf;
				//int lightId = scene->lightPowerDistribution->sampleDiscrete(rands[0],&lightPdf,&u0);
				//rands[0] = u0;
				u0 = randGen.genRand();
				int lightId = min((int)(u0*scene->lights.size()), scene->lights.size() - 1);
				lightPdf = 1.f / (scene->lights.size());
				Light* light = scene->lights[lightId];

				Ray isect_ray = { its_pt,ray.d,ray.medium };
				film[index ].Ld += beta * sampleDirectLighting(&si, bsdf, isect_ray, light, &randGen, true) / lightPdf;
				//P3 lightDir;
				//Float pdf1 = 0;
				//VisibilityTester vis;
				//Spectrum le = scene->lights[lightId]->sampleLightPt(its_pt, &lightDir, rands, &pdf1, &vis) / lightPdf;

				//if (pdf1 > 0 && !le.IsBlack())
				//{
				//	Spectrum fl = bsdf->f(-lightDir, -ray.d,BxDFType(BSDF_ALL^BSDF_SPECULAR));
				//	if (!fl.IsBlack() && vis.unoccluded(scene,id))
				//	{

				//		film[index].Ld += beta * le*nl.dot(-lightDir)*fl / (pdf1);
				//	}
				//}

				int nDiffuse = bsdf->getNumComponent(BxDFType(
					BSDF_DIFFUSE | BSDF_REFLECTION |
					BSDF_TRANSMISSION));
				int nGlossy = bsdf->getNumComponent(BxDFType(
					BSDF_GLOSSY | BSDF_REFLECTION |
					BSDF_TRANSMISSION));
				//fprintf(stderr, "nDiff=%d,nGlossy=%d\n", nDiffuse, nGlossy);
				Float pDiff = 0;
				if (nDiffuse > 0 || nGlossy > 0)
					pDiff = (Float)nDiffuse / (nDiffuse + nGlossy);
				Float pOthers = 1 - pDiff;
				bool sampleDiffuse = pDiff > eps && randGen.genRand() < pDiff;
				//bool sampleDiffuse = nDiffuse > 0;
				if (!sampleDiffuse&&depth < maxDepth - 1) 
				{
					//sample specular or glossy and keep tracing
					Float rands1[] = { randGen.genRand(),randGen.genRand() };
					//Float rands1[] = { (Float)rand() / RAND_MAX ,(Float)rand() / RAND_MAX };
					Float pdf;
					P3 nextDir;
					BxDFType sampledType;
					Spectrum f = bsdf->ImportanceSample_f(-ray.d, &nextDir, rands1, &pdf, BxDFType(BSDF_ALL^BSDF_DIFFUSE), &sampledType);
					if (pdf == 0.f || f.IsBlack()) break;
					isSpec = (sampledType & (BSDF_SPECULAR)) != 0;
					//trace next direction
					if ((sampledType&BSDF_SPECULAR) == 0) //not specular
						f = f * abs(si.n.dot(nextDir)) / pdf;
					beta *= f / (pOthers);
					//Russian roulette
					if (beta.y() < 0.25f)
					{
						Float prob = beta.y();
						if (randGen.genRand() > prob) break;
						beta /= prob;
					}
					ray.o = its_pt+nextDir*eps, ray.d = nextDir;
					MediumInterface mi = static_cast<Object*>(si.pObj)->mi;
					ray.medium = mi.getMedium(si.n, nextDir);
					continue;
				}


				if (sampleDiffuse )
				{
					beta /= pDiff;
	
					VisiblePt *vp = new VisiblePt(its_pt, -ray.d, bsdf,beta, index);
					vispts[thread_id].push_back(vp);
					break;
				}
				
			}
		}
	}

	for (int k = 0; k < num_procs; k++)
		visible_pts.insert(visible_pts.end(), vispts[k].begin(), vispts[k].end());
	delete[] vispts;
	auto end = std::chrono::system_clock::now();
	auto dif = end - start;
	fprintf(stdout, "visible points generated in : %lld(msec)\n", std::chrono::duration_cast<std::chrono::milliseconds>(dif).count());
}
void SPPM::gen_photon_path(int iter, int n_sample)
{

#pragma omp parallel for schedule(dynamic, 1)
	for (int j = 0; j < n_sample; j++)
	{
		long long seed = (long long)iter * (long long)n_sample + (long long)j;
		HaltonGenerator randGen(seed);
		//to be implemented: choose light. 
		Float rands[] = { randGen.genRand(),randGen.genRand() };
		Ray ray;
		P3 norm;
		Float pdfPos=0, pdfDir=0;
		Spectrum beta;
		//choose one light to sample based on light power distribution
		Float u0, lightPdf;
		u0 = randGen.genRand();
		int lightId = min((int)(u0*scene->lights.size()), scene->lights.size() - 1);
		lightPdf = 1.f / (scene->lights.size());

		Spectrum Le = scene->lights[lightId]->sampleLightOut(&ray, &norm, &pdfPos, &pdfDir, rands);
		if (pdfPos == 0 || pdfDir == 0 || Le.IsBlack()) continue;
		beta = abs(norm.dot(ray.d))*Le / (lightPdf*pdfPos*pdfDir);
		if (beta.IsBlack()) continue;

		for (int depth = 0; depth < maxDepth; ++depth)
		{
			Float t;
			SurfaceInteraction si;
			bool sampleMedia = false;
			bool intersect = scene->nearest_intersect(ray, &t, &si);
			if (!intersect) break;
			Object* obj = static_cast<Object*>(si.pObj);
			MediumInterface mi = obj->mi;
			Material* material = obj->material;
			P3 its_pt = ray.get(t);
			
			if (enableMedia&&ray.medium)
			{
				Float sampled_t;
				beta *= ray.medium->sampleMediumInteraction(ray, t, &sampled_t, &randGen);
				if (beta.IsBlack()) break;
				sampleMedia = sampled_t < t;
				its_pt = ray.get(sampled_t);
			}
			if (sampleMedia)
			{
				Float pdf;
				Float rands[] = { randGen.genRand(),randGen.genRand() };
				ray.d = ray.medium->sample_inScatterDir(its_pt, -ray.d, &pdf, rands);
				ray.o = its_pt;
				continue;
			}
			if (depth > 0)
			{

				// strictly speaking, we should use #pragma omp critical here.
				// it usually works without an artifact due to the fact that photons are 
				// rarely accumulated to the same measurement points at the same time (especially with QMC).
				// it is also significantly faster.
				int ansPt[3];
				if (toGrid(its_pt, ansPt))
				{
					auto list = hash_grid[hash(ansPt[0], ansPt[1], ansPt[2])];
					for (auto itr = list.begin(); itr != list.end(); itr++)
					{
						VisiblePt *vp = (*itr);
						SPPMPixel &pixel = film[vp->id];
						P3 d = vp->p - its_pt;

						if ((d.len2() <= pixel.r2)&&(vp->bsdf->n.dot(si.n)>0))
						{
							Spectrum Phi = beta * vp->bsdf->f(vp->view, -ray.d);
							pixel.Phi += Phi;
							++pixel.M;
						}
					}
				}
				//kd_tree->update(beta, ray.d, nl, its_pt);
			}
			//sample new photon ray direction
			if (!material) //at invisible boundary
			{
				ray.o = ray.get(t + eps);
				ray.medium = mi.getMedium(si.n, ray.d);
				//--depth;
				continue;
			}
			shared_ptr<BSDF> bsdf = material->createBSDF(&si, TransportMode::Radiance);
			Float pdf;
			Float rands1[] = { randGen.genRand(),randGen.genRand() };
			//Float rands1[] = { (Float)rand() / RAND_MAX,(Float)rand() / RAND_MAX };
			P3 dir;
			BxDFType sampledType;
			Spectrum f = bsdf->ImportanceSample_f(-ray.d,&dir, rands1, &pdf,BSDF_ALL,&sampledType);
			if (pdf == 0.f||f.IsBlack()) break;
			Spectrum bnew = beta * f*abs(si.n.dot(dir)) / pdf;
			//Russian roulette
			Float q = max(0, 1 - bnew.y() / beta.y());
			if (randGen.genRand() < q) break;
			beta = bnew / (1 - q);
			ray.o = its_pt, ray.d = dir;
		}

	}

}
void SPPM::save_img(char* name, int iter, int n_sample)
{
	int w = cam->w, h = cam->h;
	n_pixels = w * h;
	P3 *img = new P3[n_pixels];
	Float rgb[3];
	//density estimate
	for (int i = 0; i < n_pixels; ++i)
	{
		//direct illum
		//Spectrum col = film[i].ld_cnt==0?Spectrum(0):film[i].Ld/ film[i].ld_cnt;
		Spectrum col = film[i].Ld / iter;
		if (n_sample != 0)
			col += film[i].tau / (iter*n_sample*Pi*film[i].r2); //indirect illum
		col.ToRGB(rgb);
		img[i].x = rgb[0], img[i].y = rgb[1], img[i].z = rgb[2];
	}
	save_to_bmp(name, w, h, &img[0].x, 2.2);
	delete[] img;
}

void SPPM::render(string _outputname, string config)
{
	int w = cam->w, h = cam->h, n_pixels = w * h;
	outputname = _outputname;
	int n_iter, n_sample;
	if (config == "")
	{
		n_iter = 100;
		n_sample = 10000;
	}
	else
	{
		sscanf(config.c_str(), "%d,%d", &n_iter, &n_sample);
	}
	fprintf(stdout, "SPPM rendering with %d iters and %d samples per iter.\n", n_iter, n_sample);

	hash_grid.resize(n_pixels);
	hash_grid.shrink_to_fit();
	film = new SPPMPixel[n_pixels];
	for (int i = 0; i < n_pixels; ++i)
	{
		film[i].r2 = INITRADIUS * INITRADIUS;
	}
	//visible_pts.reserve(n_pixels);
	//gen_camera_path(n_sample);
	//fprintf(stderr, "num of vps:%d\n", visible_pts.size());
	//build_hash_grid(w, h);

	for (int i = 0; i < n_iter; i++)
	{
		fprintf(stderr, "iter %d of %d\n", i + 1, n_iter);
		long long seed = (long long)i*(long long)n_sample+i;
		hpbox.reset();
		//srand(i);
		gen_camera_path(seed);
		//fprintf(stderr, "num of vps:%d\n", visible_pts.size());
		if (visible_pts.size() == 0)
		{
			if (i != 0 && i % 1000 == 0)
			{
				char filename[100];
				sprintf(filename, "out/%s_%03d.bmp", outputname.c_str(), i);
				save_img(filename, i + 1, n_sample);
			}
			continue;
		}
		auto start = chrono::system_clock::now();
		build_hash_grid(w, h);
		auto end = std::chrono::system_clock::now();
		auto dif = end - start;
		fprintf(stdout, "hash table built in : %lld(msec)\n", std::chrono::duration_cast<std::chrono::milliseconds>(dif).count());
		gen_photon_path(i, n_sample);
		//after one iteration: update pixel info
		for (auto itr = visible_pts.begin(); itr != visible_pts.end(); itr++)
		{
			VisiblePt *vp = (*itr);
			SPPMPixel &pixel = film[vp->id];
			if (pixel.M > 0)
			{
				Float Nnew = pixel.N + GAMMA * (Float)pixel.M;
				Float R2new = pixel.r2*Nnew / (pixel.N + pixel.M);
				pixel.tau = (pixel.tau + vp->beta*pixel.Phi)*R2new / pixel.r2;
				pixel.N = Nnew;
				pixel.r2 = R2new;
				//reset M and Phi
				pixel.M = 0;
				pixel.Phi = Spectrum();
			}
			//reset vp

		}
		//reset vp
		for (int j = 0; j < n_pixels; ++j)
		{
			hash_grid[j].resize(0);
			hash_grid[j].shrink_to_fit();
		}
		for (auto itr = visible_pts.begin(); itr != visible_pts.end(); itr++) delete (*itr);
		visible_pts.resize(0);
		visible_pts.shrink_to_fit();
		if (i != 0 && i % 100 == 0)
		{
			char filename[100];
			sprintf(filename, "out/%s_%03d.bmp", outputname.c_str(), i);
			save_img(filename, i + 1, n_sample);
		}
	}

	char filename[100];
	sprintf(filename, "out/%s.bmp", outputname.c_str());
	save_img(filename, n_iter + 1, n_sample);
}

#endif // !M_PPM

